import numpy as np
from numpy.typing import NDArray
from typing import List

from hj_dataset_devkit import BoundingBox

from model_frame import DetectObject

class BBox:
    def __init__(self, x: float=None, y: float=None, z: float=None,
                 h: float=None, w: float=None, l: float=None,
                 o: float=None, s: float=None) -> None:
        self.x = x  # center x
        self.y = y  # center y
        self.z = z  # center z
        self.h = h  # height (along z)
        self.w = w  # width (along y)
        self.l = l  # length (along x)
        self.o = o  # orientation
        self.s = s  # detection score

    def to_array(self, mmorder=False) -> NDArray:
        if self.s is None:
            if mmorder:
                return np.array([self.x, self.y, self.z, self.w, self.l, self.h, self.o])
            else:
                return np.array([self.x, self.y, self.z, self.o, self.l, self.w, self.h])
        else:
            if mmorder:
                return np.array([self.x, self.y, self.z, self.w, self.l, self.h, self.o, self.s])
            else:
                return np.array([self.x, self.y, self.z, self.o, self.l, self.w, self.h, self.s])

    def to_mmarray(self) -> NDArray:
        return np.array([self.x, self.y, self.z, self.w, self.l, self.h, self.o])

    def to_corners2d(self) -> List[list]:
        """ the coordinates for bottom corners
        """
        bottom_center = np.array([self.x, self.y, self.z - self.h / 2])
        cos, sin = np.cos(self.o), np.sin(self.o)
        pc0 = np.array([self.x + cos * self.l / 2 + sin * self.w / 2,
                        self.y + sin * self.l / 2 - cos * self.w / 2,
                        self.z - self.h / 2])
        pc1 = np.array([self.x + cos * self.l / 2 - sin * self.w / 2,
                        self.y + sin * self.l / 2 + cos * self.w / 2,
                        self.z - self.h / 2])
        pc2 = 2 * bottom_center - pc0
        pc3 = 2 * bottom_center - pc1
        return [pc0.tolist(), pc1.tolist(), pc2.tolist(), pc3.tolist()]

    def to_boundingbox(self) -> BoundingBox:
        return BoundingBox([self.x, self.y, self.z],
                           [self.l, self.w, self.h],
                           [0, 0, self.o])

    @classmethod
    def from_detect_obj(cls, obj: DetectObject):
        return cls(obj.center.x, obj.center.y, obj.center.z,
                   obj.size.z, obj.size.y, obj.size.x,
                   obj.rotation.z, obj.score)

    @classmethod
    def from_array(cls, data: NDArray):
        bbox = cls()
        bbox.x, bbox.y, bbox.z, bbox.o, bbox.l, bbox.w, bbox.h = data[:7]
        if len(data) == 8:
            bbox.s = data[-1]
        return bbox

    @classmethod
    def merge_boxes(cls, box_list: List['BBox']) -> NDArray:
        if len(box_list) == 0:
            boxes = np.zeros((0, 7), dtype=np.float32)
        else:
            boxes = np.stack([b.to_mmarray() for b in box_list], 0)
        assert boxes.shape[1] == 7
        heading = -boxes[:, 6] - 0.5 * np.pi
        while (heading < -np.pi).any():
            heading[heading < -np.pi] += 2 * np.pi
        while (heading > np.pi).any():
            heading[heading > np.pi] -= 2 * np.pi
        boxes[:, 6] = heading
        return boxes
